#include <stdio.h>

#include <cuda.h>
#include <omp.h>
#include <stdlib.h>

#include <cufft.h>
#include <cuComplex.h>
#include <cufft.h>

#include "cush_alloc.h"
#include "cuda_util.h"

/**
 *
 * Helper functions
 *
 */ 

size_t cush_t_size(cush_t type){
    int size = 0;
    switch(type){
        case INT:
            size = sizeof(int);
            break;
        case FLOAT:
            size = sizeof(float);
            break;
        case DOUBLE:
            size = sizeof(double);
            break;
        case CUFFT_COMPLEX:
            size = sizeof(cufftComplex);
        default:
            // intentionally left blank
            break;
    }

    return size;
}

//////////////////////////////////////////////////////////////

void cush_dp_init(cush_dbuff_t* dp, cush_t type){
    dp->type = type;
    dp->type_size = cush_t_size(type);
}

void cush_dp_init(cush_dbuff_t* dp, cush_t type, size_t n_devs){
    cush_dp_init(dp, type);

    dp->ndevices = n_devs;
    dp->id      = (int*)malloc(n_devs*sizeof(int));
    dp->size    = (size_t*)malloc(n_devs*sizeof(size_t));
    dp->offset  = (size_t*)malloc(n_devs*sizeof(size_t));
    dp->ptr     = (void***)malloc(n_devs*sizeof(void***));
    dp->stream  = (cudaStream_t*)malloc(n_devs*sizeof(cudaStream_t));

    for(int i = 0; i < n_devs; i++){
        cudaStreamCreate(&(dp->stream[i]));
    }
}

/**
 *
 * Sets up a pseudo-uniform distribution allocation sizes.
 *
 */
void cush_unif_size_dist(cush_dbuff_t* dbuff, size_t host_size){

	size_t base_size = host_size / dbuff->ndevices;
	size_t size_remaining = host_size % base_size;

	int i = 0;
    size_t offset = 0;
	for(i = 0; i < dbuff->ndevices; i++){
		dbuff->size[i] = base_size;
		if(size_remaining > 0){
			dbuff->size += dbuff->type_size;
			size_remaining -= dbuff->type_size;
		}
		dbuff->offset[i] = offset;
		offset += dbuff->size[i];
	}
}

/**
 *
 * Allocates memory of size dbuff->size[i] to each device i with 
 * device pointers stored in dbuff->ptr[i].
 *
 * The size of allocations are done by helper functions such as "cush_unif_size_dist()"
 * or defined by the user.
 *
 */
void cush_alloc(cush_dbuff_t* dbuff){
	
	int i = 0;
	#pragma omp parallel for num_threads(dbuff->ndevices)
	for(i = 0; i < dbuff->ndevices; i++){
	    int th_id = omp_get_thread_num();
        //printf("Cush allocing %.4fMB on dev %d\n",dbuff->size[i]/1048576.0, dbuff->id[i]);
        //int th_id = dbuff->id[i];
		cudaSetDevice(dbuff->id[i]);
		cudaMalloc(&(dbuff->ptr[i]), dbuff->size[i]);
		//cuMemAlloc(&(dbuff->ptr[i]),  dbuff->size[i]);
		checkCUDAError("cush malloc failed",th_id);
	}
}

// Nice syntactic sugar
void cush_dbuff_create(cush_dbuff_t* dbuff, cush_t type, size_t n_devs, size_t host_buff_size, cush_dist_t dist, int* dev_ids){
    cush_dp_init(dbuff, type, n_devs);
    
    switch(dist){
    case UNIF:
            cush_unif_size_dist(dbuff, host_buff_size);
            break;
        default:
            // intentionally left blank
            break;
    }
    
    for(int i = 0; i < n_devs; i++){
		dbuff->id[i] = dev_ids[i];
    }

    cush_alloc(dbuff);
}

/* vim: set sw=4 sts=4 et foldmethod=syntax syntax=c : */
